import tensorflow as tf
from tensorflow.keras.callbacks import ModelCheckpoint
from tensorflow.keras.callbacks import TensorBoard
from pathlib import Path
from datetime import datetime
from absl import app, flags
from absl.flags import FLAGS
from termcolor import colored

from models.yolonet import *
from tools.metric import Yolo_Precision, Yolo_Recall
from tools.dataset import Dataset
from tools.loss import YoloLoss

INFO = colored('[ INFO  ]', 'blue')
ERROR = colored('[ ERROR ]', 'red')
NOTE = colored('[ NOTE ]', 'green')

flags.DEFINE_string('dataset', 'coco', 'type of dataset')
flags.DEFINE_string('save_dir', './log/', 'path to save model')
flags.DEFINE_string('pre_ckpt', None, 'path to pretrained weight')
flags.DEFINE_string('model_type', 'yolo_mobilev1', 'type of model')

flags.DEFINE_enum('mode', 'fit', ['fit', 'eager'],
                  'fit: model.fit, '
                  'eager: custom GradientTape')

flags.DEFINE_integer('img_h', 224, 'image height')
flags.DEFINE_integer('img_w', 320, 'image width')
flags.DEFINE_integer('epochs', 40, 'number of epochs')
flags.DEFINE_integer('batch_size', 64, 'batch size')
flags.DEFINE_integer('class_num', 80, 'number of classes in the model')

flags.DEFINE_float('alpha', 0.5, 'alpha of mobilenet')
flags.DEFINE_float('learning_rate', 1e-3, 'learning rate')
flags.DEFINE_float('obj_thresh', 0.5, 'object thresh')


def main(_argv):
    gpus = tf.config.experimental.list_physical_devices(device_type='GPU')
    tf.config.experimental.set_visible_devices(devices=gpus[0], device_type='GPU')
    tf.config.experimental.set_memory_growth(device=gpus[0], enable=True)

    # Build path
    log_dir = (Path(FLAGS.save_dir) / datetime.strftime(datetime.now(), '%Y%m%d-%H%M%S'))  # type: Path
    ckpt = log_dir / 'yolo_model.h5'
    if not log_dir.exists():
        log_dir.mkdir(parents=True)

    # Build dataset
    dataset_dir = 'data/{}/{}_img_ann.npy'.format(FLAGS.dataset.lower(), FLAGS.dataset.lower())
    anchor_dir = 'data/{}/{}_anchor.npy'.format(FLAGS.dataset.lower(), FLAGS.dataset.lower())
    ds = Dataset(dataset_dir, anchor_dir, FLAGS.batch_size, FLAGS.class_num, (FLAGS.img_h, FLAGS.img_w))
    ds.set_dataset()

    # Build network
    network = eval(FLAGS.model_type)
    yolo_model, yolo_model_warpper = network([FLAGS.img_h, FLAGS.img_w, 3], ds.n_anchor, FLAGS.class_num, alpha=FLAGS.alpha)

    yolo_model.summary()

    if FLAGS.pre_ckpt != None and FLAGS.pre_ckpt != 'None' and FLAGS.pre_ckpt != '':
        if 'h5' in FLAGS.pre_ckpt:
            yolo_model_warpper.load_weights(str(FLAGS.pre_ckpt))
            print(INFO, f' Load CKPT {str(FLAGS.pre_ckpt)}')
        else:
            print(ERROR, ' Pre CKPT path is unvalid')

    train_model = yolo_model_warpper

    filepath = str(log_dir) + '/yolo_model_{epoch:02d}_{val_loss:.2f}.h5'
    checkpoint = ModelCheckpoint(filepath=filepath, monitor='val_loss', mode='auto', save_best_only='True', save_weights_only='True')

    # reduce_lr = keras.callbacks.ReduceLROnPlateau(monitor='val_loss', factor=0.5,
    #                                               patience=2, verbose=1, mode='auto',
    #                                               min_delta=0.0001, cooldown=2, min_lr=0)
    reduce_lr = keras.callbacks.LearningRateScheduler(
        keras.experimental.CosineDecayRestarts(FLAGS.learning_rate, 10), verbose=1)

    train_model.compile(
        # keras.optimizers.Adadelta(),
        keras.optimizers.Adam(lr=FLAGS.learning_rate),
        loss=[YoloLoss(ds.anchors[mask], classes=FLAGS.class_num, batch_size=FLAGS.batch_size)
              for mask in ds.anchor_masks],
        metrics=[Yolo_Precision(FLAGS.obj_thresh, name='p'), Yolo_Recall(FLAGS.obj_thresh, name='r')])

    """ Callbacks """
    cbs = [TensorBoard(str(log_dir), update_freq='batch', profile_batch=3),
           reduce_lr,
           checkpoint]

    # Training
    try:
        train_model.fit(ds.train_dataset,
                        epochs=FLAGS.epochs,
                        steps_per_epoch=ds.train_epoch_step,
                        callbacks=cbs,
                        validation_data=ds.valid_dataset,
                        validation_steps=ds.valid_epoch_step,
                        use_multiprocessing=True)
    except KeyboardInterrupt as e:
        pass

    keras.models.save_model(yolo_model, str(ckpt))
    print()
    print(INFO, f' Save Model as {str(ckpt)}')


if __name__ == "__main__":
    try:
        app.run(main)
    except SystemExit:
        pass
